package musicalgo;

import ai.neuralnet.ClassificationDataSet;
import ai.neuralnet.NNTrainer;
import ai.neuralnet.NeuralNetwork;
import musicalgo.data.DataIO;
import shared.MathHelper;

import java.util.ArrayList;
import java.util.List;

public class VotingNNAlgorithm extends NeuralNetworkAlgorithm
{

    private int families = 10;

    double[] voteWeights;

    private int categories = 5;

    private List<NeuralNetwork> votersList;

    public VotingNNAlgorithm(NNTrainer trainer)
    {
        super(trainer);
        votersList = new ArrayList<NeuralNetwork>(families);
    }


    @Override
    protected void train()
    {
        double percent = trainingPercentage / families;


        int dataSize = trainingData.size();

        for (int i = 0; i < families; i++)
        {
            //normal size directly from percentage divided over families
            int size = (int) (percent * dataSize);
            //actual percentage including overlap
            double overPercent = MathHelper.clamp(1.2 * percent, percent, trainingPercentage);
            //offset from first song in list to be fist song in subset
            int offset = (dataSetOffset + i * size) % dataSize;

            trainer.setTrainingSet(getTrainingSet(overPercent, offset));
            System.out.println("Training set loaded with " + trainer.getTrainingSet().size() + " Songs");
            //same parameters gets complementary set from training set method
            trainer.setValidationSet(getValidationSet(overPercent, offset));
            System.out.println("Validation set loaded with " + trainer.getValidationSet().size() + " Songs");


            System.out.println("Training Voter Family " + (i + 1) + " training offset: " + offset);
            trainer.trainNetworks();

            List<NeuralNetwork> networks = trainer.getNetworks();
            System.out.println("Family " + (i + 1) + " Trained:");
            for (int k = 0; k < networks.size(); k++)
            {
                System.out.println("Voter " + (k + 1) + " Validation Accuracy: " + trainer.getValidationSet().getAccuracy(networks.get(k))[0]);
            }
            votersList.addAll(trainer.getNetworks());
        }
        System.out.println("Computing Voter Weights");
        ClassificationDataSet allSongs = getTrainingSubset(0, dataSize);
        trainer.setValidationSet(allSongs);
        trainer.setNetworks(votersList);

        voteWeights = trainer.validationAccuracies();
        double sum = 0;
        for (int i = 0; i < voteWeights.length; i++)
        {
            System.out.println("Total Set Voter " + (i + 1) + " Accuracy : " + voteWeights[i]);
            sum += voteWeights[i];
        }
        System.out.println("Average Voter Accuracy: " + sum / voteWeights.length);


        int[] pred = votePredictions(allSongs);

        double acc = allSongs.getAccuracy(pred);

        System.out.println("Accuracy from Votes: " + acc);


    }

    private int[] votePredictions(ClassificationDataSet set)
    {


        double[][] predictions = new double[set.size()][categories];

        System.out.println("Voting");
        for (int i = 0; i < families; i++)
        {
            NeuralNetwork voter = votersList.get(i);
            double w = voteWeights[i];
            int[] vpred = set.getPredictions(voter);
            for (int j = 0; j < vpred.length; j++)
            {
                predictions[j][vpred[j]] += w;
            }


        }

        System.out.println("Counting Votes");
        int[] finalPredictions = new int[set.size()];
        for (int i = 0; i < finalPredictions.length; i++)
        {
            double max = predictions[i][0];
            int maxI = 0;
            for (int j = 1; j < categories; j++)
            {
                if (predictions[i][j] > max)
                {
                    max = predictions[i][j];
                    maxI = j;
                }
            }

            finalPredictions[i] = maxI;
        }

        return finalPredictions;
    }

    @Override
    protected void test()
    {

        System.out.println("Classifying songs ...");


        ClassificationDataSet testSet = getTestingSet();

        int[] predictions = votePredictions(testSet);

        System.out.println("Writing output");

        DataIO.writePredictions("testy.csv", predictions);
    }

    public int getFamilies()
    {
        return families;
    }

    public void setFamilies(int families)
    {
        this.families = families;
    }
}
